import torch

from typing import List
from torch import Tensor


def calculate_cnn_output_size(input_size: int,
                              channels: List[int],
                              padding: List[int],
                              kernel_size: List[int],
                              stride: List[int]) -> int:
    for p, k, s in zip(padding, kernel_size, stride):
        input_size = 1 + (input_size + 2 * p - k) // s
    return input_size * channels[-1]
    

def init_model_weights(w: Tensor):
    if isinstance(w, torch.nn.Linear):
        torch.nn.init.xavier_normal_(w.weight)
        torch.nn.init.constant_(w.bias, val=0.0)


def init_transformer_weights(transformer: torch.nn.TransformerEncoder):
    for name, param in transformer.named_parameters():
        if 'weight' in name and param.data.dim() == 2:
            torch.nn.init.xavier_normal_(param)